import cv2
import numpy as np
from datetime import datetime

# -------------------------- 1. 配置参数（需根据图像修改）--------------------------
IMAGE_PATH = "butterfly_with_ruler.jpeg"  # 带圆弧弯曲的蝴蝶图像路径
BLUR_KERNEL = (3, 3)              # 高斯模糊核（去噪）
THRESHOLD_VALUE = 130             # 二值化阈值（调整至翅膀轮廓清晰）
EXPORT_TXT_PATH = "butterfly_arc_result4.txt"  # 结果导出路径

# 全局变量：存储鼠标点击的点（圆弧上的点）
arc_points = []
# 存储拟合的椭圆参数：(中心(x,y), (长轴,短轴), 旋转角度)
ellipse_params = None
# 最终计算的圆心角（度）
arc_angle = 0.0


# -------------------------- 2. 鼠标交互函数（选点+绘制）--------------------------
def mouse_callback(event, x, y, flags, param):
    global arc_points, ellipse_params, arc_angle
    img_copy = param  # 传入图像副本，避免原图被反复修改

    # 左键点击：记录圆弧上的点
    if event == cv2.EVENT_LBUTTONDOWN:
        arc_points.append((x, y))
        # 在图像上绘制选中的点（红色，大小5）
        cv2.circle(img_copy, (x, y), 5, (0, 0, 255), -1)
        # 标记点的序号
        cv2.putText(img_copy, str(len(arc_points)), (x+10, y-10), 
                    cv2.FONT_HERSHEY_SIMPLEX, 0.6, (255, 0, 0), 2)
        print(f"已选点{len(arc_points)}: ({x}, {y})")

        # 选够3个点：拟合椭圆并计算圆心角
        if len(arc_points) >= 3:
            # 转换为OpenCV需要的点格式（(n,1,2)的uint8数组）
            pts = np.array(arc_points, dtype=np.int32).reshape((-1, 1, 2))
            try:
                # 拟合椭圆（至少需要5个点更稳定，3个点也可尝试）
                ellipse_params = cv2.fitEllipse(pts)
                # 绘制拟合的椭圆（绿色）
                cv2.ellipse(img_copy, ellipse_params, (0, 255, 0), 2)
                
                # 计算圆心角（取首尾点与椭圆中心的夹角）
                if len(arc_points) >= 2:
                    center = (int(ellipse_params[0][0]), int(ellipse_params[0][1]))
                    pt_start = arc_points[0]
                    pt_end = arc_points[-1]
                    arc_angle = calculate_arc_angle(center, pt_start, pt_end)
                    # 显示圆心角
                    cv2.putText(img_copy, f"圆弧角: {arc_angle:.1f}°", 
                                (50, 50), cv2.FONT_HERSHEY_SIMPLEX, 1.2, (0, 0, 255), 3)
                    print(f"计算完成！圆弧对应角度: {arc_angle:.1f}°")
            except:
                print("提示：选点不足或分布不合理，无法拟合椭圆！")


# -------------------------- 3. 圆心角计算函数 --------------------------
def calculate_arc_angle(center, pt_start, pt_end):
    """计算椭圆中心到两点的夹角（圆心角）"""
    # 向量1：中心→起点
    vec1 = (pt_start[0] - center[0], pt_start[1] - center[1])
    # 向量2：中心→终点
    vec2 = (pt_end[0] - center[0], pt_end[1] - center[1])
    
    # 点积公式计算夹角（弧度转角度）
    dot_product = vec1[0]*vec2[0] + vec1[1]*vec2[1]
    norm_vec1 = np.linalg.norm(vec1)
    norm_vec2 = np.linalg.norm(vec2)
    cos_theta = dot_product / (norm_vec1 * norm_vec2)
    cos_theta = max(min(cos_theta, 1.0), -1.0)  # 避免数值误差
    angle_rad = np.arccos(cos_theta)
    angle_deg = np.degrees(angle_rad)
    return angle_deg


# -------------------------- 4. 结果导出函数 --------------------------
def export_arc_result(angle, img_path, ellipse_info=None):
    """导出圆弧测量结果到TXT文件"""
    with open(EXPORT_TXT_PATH, "w", encoding="utf-8") as f:
        f.write(f"蝴蝶翅膀圆弧弯曲测量报告\n")
        f.write(f"测量时间：{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n")
        f.write(f"图像路径：{img_path}\n")
        f.write(f"圆弧对应圆心角：{angle:.2f} 度\n")
        if ellipse_info:
            center, axes, angle = ellipse_info
            f.write(f"拟合椭圆中心：({center[0]:.1f}, {center[1]:.1f})\n")
            f.write(f"拟合椭圆长轴：{axes[0]:.1f}, 短轴：{axes[1]:.1f}\n")
            f.write(f"拟合椭圆旋转角：{angle:.1f} 度\n")
    print(f"\n结果已导出至：{EXPORT_TXT_PATH}")


# -------------------------- 5. 主流程执行 --------------------------
if __name__ == "__main__":
    # 步骤1：图像预处理（突出翅膀轮廓）
    img = cv2.imread(IMAGE_PATH)
    if img is None:
        print(f"错误：无法读取图像 {IMAGE_PATH}")
        exit(1)
    img_gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)
    img_blur = cv2.GaussianBlur(img_gray, BLUR_KERNEL, 0)
    # 二值化（翅膀为白色，背景为黑色）
    _, img_bin = cv2.threshold(img_blur, THRESHOLD_VALUE, 255, cv2.THRESH_BINARY_INV)
    # 形态学膨胀（增强轮廓连续性）
    kernel = np.ones((2, 2), np.uint8)
    img_bin = cv2.dilate(img_bin, kernel, iterations=1)

    # 步骤2：创建图像副本用于绘制
    img_display = img.copy()
    # 显示操作提示
    cv2.putText(img_display, "操作提示：点击圆弧上的点（至少3个）", (20, 30), 
                cv2.FONT_HERSHEY_SIMPLEX, 0.8, (0, 255, 0), 2)

    # 步骤3：绑定鼠标回调函数，开启交互选点
    cv2.namedWindow("Butterfly Wing Arc Measurement", cv2.WINDOW_NORMAL)
    cv2.setMouseCallback("Butterfly Wing Arc Measurement", mouse_callback, img_display)

    # 步骤4：等待用户操作（按ESC退出，按S保存结果）
    while True:
        cv2.imshow("Butterfly Wing Arc Measurement", img_display)
        key = cv2.waitKey(1) & 0xFF
        # 按ESC键：退出程序
        if key == 27:
            break
        # 按S键：导出结果（需先拟合椭圆）
        if key == ord('s'):
            if ellipse_params is not None:
                export_arc_result(arc_angle, IMAGE_PATH, ellipse_params)
            else:
                print("提示：请先选点拟合椭圆再导出结果！")

    # 释放资源
    cv2.destroyAllWindows()